import numpy as np
import cv2 as cv


# 相机内参
P2 = np.mat([[1916.240369950428,    0,                  0],
            [0,                     1913.278508450042,  0],
            [1021.529553590090,     532.6943887717723,  1],
            [0,                     0,                  0]]).T

R0_rect = np.mat([[-1.702738287348829968e-02, -9.998550236071626074e-01, 0.000000000000000000e+00, 1.302083333333332611e-03],
                  [0.000000000000000000e+00, 0.000000000000000000e+00, -1.000000000000000000e+00, 6.315104166666665186e-01],
                  [9.998550236071626074e-01, -1.702738287348829968e-02, 0.000000000000000000e+00, 2.144531250000000444e+00],
                  [0.000000000000000000e+00, 0.000000000000000000e+00, 0.000000000000000000e+00, 1.000000000000000000e+00]])
trans_mat = P2 * R0_rect


# 投影函数
def pcd_proj(points):
    velo = np.insert(points, 3, 1, axis=1).T
    cam = trans_mat * velo
    cam = np.delete(cam, np.where(cam[2, :] < 0)[1], axis=1)
    velo = np.delete(velo, np.where(cam[2, :] < 0)[1], axis=1)
    cam[:2] /= cam[2, :]
    u, v, tmp1 = cam
    u1, v1, z1, tmp2 = velo
    return u, v, u1, v1, z1


img_data = cv.imread('tmp_img/tmp_27.png', 1)
pcd_data = np.load('tmp_pcd/tmp_9.npy')
pcd_data = np.mat(pcd_data)
u0, v0, u_1, v_1, z_1 = pcd_proj(pcd_data)
for i_u in range(u0.shape[1]):
    width = int(u0[0, i_u])
    height = int(v0[0, i_u])
    dis = np.sqrt(u_1[0, i_u] ** 2 + v_1[0, i_u] ** 2 + z_1[0, i_u] ** 2)
    if dis > 8:
        cv.circle(img_data, (width, height), radius=2, color=(255, 0, 0), thickness=1)
    elif dis > 6:
        cv.circle(img_data, (width, height), radius=2, color=(255, 255, 0), thickness=1)
    elif dis > 4:
        cv.circle(img_data, (width, height), radius=2, color=(0, 255, 0), thickness=1)
    else:
        cv.circle(img_data, (width, height), radius=2, color=(0, 255, 255), thickness=1)
show_img = cv.resize(img_data, [960, 540])
while True:
    cv.imshow('tmp', show_img)
    cv.waitKey(10)
